from __future__ import annotations

from collections import OrderedDict

from logger import logger
from model import (
    Settings,
    Action,
    ReloadAction,
    VisitAction,
    OffloadAction,
    FinAction,
    Request,
)
from utils.intervals import IntervalSet


class Validator:
    def __init__(
        self,
        settings: Settings,
        requests: list[Request],
        actions: list[Action] | None,
        verbose: bool = False,
    ) -> None:
        self.settings = settings
        self.requests = requests
        self.actions = actions
        self.verbose = verbose
        self.requests_by_start: OrderedDict[int, list[Request]] = OrderedDict()
        for req in self.requests:
            if req.start not in self.requests_by_start:
                self.requests_by_start[req.start] = []
            self.requests_by_start[req.start].append(req)

    def validate(self) -> bool:
        try:
            logger.info("Starting validation ⚡️")
            self.validate_requests()
            if self.actions is not None:
                logger.info("------------------------------")
                self.validate_actions()
            logger.info("Validation completed successfully ✅")
            return True
        except ValueError as e:
            logger.error("Validation failed ❌")
            if self.verbose:
                logger.exception(e)
            else:
                logger.error(e)
                logger.error(
                    "Rerun with verbose mode (`-v` or `--verbose`) for more details."
                )
            return False

    def validate_requests(self) -> None:
        """
        Rules:
        - Chronology:
          - `request`s should be ordered by `request.start`.
        - Integrity:
          - For those with same `request.start`, their `request.time` must be the same.
        - Memory Bound:
          - Every `request` should not exceed memory limit.
          - For those with same `request.start`, sum of their memory usage should not exceed memory limit.
            - Note that overlapping memory regions should be counted only once.
        """
        # Chronology check
        logger.info("Validating Requests - chronology...")
        for i in range(1, len(self.requests)):
            if self.requests[i].start < self.requests[i - 1].start:
                raise ValueError(
                    f"requests are not ordered by start: {self.requests[i - 1]} comes before {self.requests[i]}"
                )

        # Integrity check
        logger.info("Validating Requests - integrity...")
        for start, reqs in self.requests_by_start.items():
            times = set(req.time for req in reqs)
            if len(times) != 1:
                raise ValueError(
                    f"requests starting at {start} have different times: {times}"
                )

        # Memory Bound check
        logger.info("Validating Requests - memory bound...")
        for start, reqs in self.requests_by_start.items():
            if self.verbose:
                logger.debug(f"-> validating requests starting at {start}: {reqs}")

            # Calculate total memory usage
            loaded_intervals = IntervalSet()
            for req in reqs:
                loaded_intervals.add(req.addr, req.addr + req.size)
            total_memory = loaded_intervals.total()

            if self.verbose:
                logger.debug(f"--> merged intervals: {loaded_intervals}")

            if self.verbose:
                logger.debug(
                    f"--> total memory usage: {total_memory} (limit: {self.settings.M})"
                )
            if total_memory > self.settings.M:
                raise ValueError(
                    f"total memory usage {total_memory} exceeds limit {self.settings.M} for requests starting at {start}"
                )
        logger.info("Validated Requests")

    def validate_actions(self) -> None:
        """
        Rules:
        - Chronology:
          - `action`s should be ordered by `action.time`.
        - Integrity:
          - `FinAction` must be the last action in the sequence, and only one `FinAction` should exist.
          - `VisitAction` should refer to valid request indices one by one in ascending order (i.e. `0`, `1`, `2`, ...).
          - The time of `VisitAction` (`action.time`) should be greater than or equal to the `start` time of the corresponding `request` (`request.start`).
        - Memory Access & Bound:
          - The memory region accessed by `VisitAction` should be within the loaded memory (via `ReloadAction`).
          - The loaded memory should not exceed memory limit (`ReloadAction` increases memory usage, while `OffloadAction` decreases it).
            - Note that overlapping memory regions should be counted only once.
          - If a memory region is being used by a `VisitAction`, it must not be offloaded by an `OffloadAction` before the `VisitAction` is completed.
        - Concurrency:
          - `ReloadAction` and `OffloadAction` cannot be performed concurrently.
          - IO Actions (`ReloadAction`, `OffloadAction`) and `VisitAction` can be performed concurrently.
          - If some requests have same `request.start`, their corresponding `VisitAction` must be concurrent and start at the same time.
            Otherwise, `VisitAction` must be sequential.
        - Report Correctness:
          - `FinAction` must report correct (accurate) elapsed time.
        """
        # Chronology check
        logger.info("Validating Actions - chronology...")
        if len(self.actions) < 1:
            raise ValueError("There must be at least one action")
        for i in range(1, len(self.actions)):
            if self.actions[i].time < self.actions[i - 1].time:
                raise ValueError(
                    f"Actions are not ordered by time: '{self.actions[i - 1]}' comes before '{self.actions[i]}'"
                )

        # Integrity check
        logger.info("Validating Actions - integrity...")
        if not isinstance(self.actions[-1], FinAction):
            raise ValueError("The last action must be an FinAction")
        if (count := len([a for a in self.actions if isinstance(a, FinAction)])) != 1:
            raise ValueError(f"There must be exactly one FinAction, found {count}")
        visit_actions = [a for a in self.actions if isinstance(a, VisitAction)]
        if len(visit_actions) != len(self.requests):
            raise ValueError(
                f"The number of VisitActions ({len(visit_actions)}) does not match the number of requests ({len(self.requests)})"
            )
        for i in range(0, len(visit_actions)):
            if visit_actions[i].index < 0 or visit_actions[i].index >= len(
                self.requests
            ):
                raise ValueError(
                    f"VisitAction index out of range: {visit_actions[i].index}"
                )
            if visit_actions[i].index != i:
                raise ValueError(
                    f"VisitAction indices must be in ascending order without gaps: expected {i}, found {visit_actions[i].index}"
                )
            if visit_actions[i].time < visit_actions[i].request.start:
                raise ValueError(
                    f"VisitAction time {visit_actions[i].time} is less than corresponding Request start time {visit_actions[i].request.start}"
                )

        # Memory Access & Bound check
        logger.info("Validating Actions - memory access & bound...")
        loaded_intervals = IntervalSet()
        # Track active VisitActions (those that have started but not yet completed)
        active_visits: list[VisitAction] = []

        for action in self.actions[:-1]:  # Exclude FinAction
            # Remove completed VisitActions from active list
            active_visits = [
                v for v in active_visits if action.time < v.time + v.cost()
            ]

            if self.verbose:
                logger.debug(f"-> validating action: '{action}'")
            if isinstance(action, ReloadAction):
                difference = loaded_intervals.difference(
                    action.addr, action.addr + action.size
                )
                if (action.addr, action.addr + action.size) not in difference:
                    logger.warning(
                        f"ReloadAction '{action}' overlaps with already loaded memory regions. Load {difference} is enough."
                    )
                loaded_intervals.add(action.addr, action.addr + action.size)
                total_loaded_memory = loaded_intervals.total()
                if self.verbose:
                    logger.debug(
                        f"--> total loaded memory: {total_loaded_memory} (limit: {self.settings.M})"
                    )
                if total_loaded_memory > self.settings.M:
                    raise ValueError(
                        f"total loaded memory {total_loaded_memory} exceeds limit {self.settings.M} after {action}"
                    )
            elif isinstance(action, OffloadAction):
                # Check if offload conflicts with any active VisitAction
                offload_start = action.addr
                offload_end = action.addr + action.size
                for visit in active_visits:
                    visit_start = visit.request.addr
                    visit_end = visit.request.addr + visit.request.size
                    # Check if memory regions overlap
                    if offload_start < visit_end and offload_end > visit_start:
                        raise ValueError(
                            f"OffloadAction '{action}' conflicts with active VisitAction '{visit}'. "
                            f"Memory region [{offload_start}, {offload_end}) overlaps with "
                            f"[{visit_start}, {visit_end}) which is in use until time {visit.time + visit.cost()}."
                        )

                difference = loaded_intervals.difference(
                    action.addr, action.addr + action.size
                )
                if (
                    difference
                    and (action.addr, action.addr + action.size) not in difference
                ):
                    logger.warning(
                        f"OffloadAction '{action}' includes regions not currently loaded. Offload {difference} is unnecessary."
                    )
                loaded_intervals.remove(action.addr, action.addr + action.size)
                total_loaded_memory = loaded_intervals.total()
                if self.verbose:
                    logger.debug(
                        f"--> total loaded memory: {total_loaded_memory} (limit: {self.settings.M})"
                    )
                if total_loaded_memory > self.settings.M:
                    raise ValueError(
                        f"total loaded memory {total_loaded_memory} exceeds limit {self.settings.M} after '{action}'"
                    )
            elif isinstance(action, VisitAction):
                if not loaded_intervals.contains(
                    action.request.addr, action.request.addr + action.request.size
                ):
                    difference = loaded_intervals.difference(
                        action.request.addr, action.request.addr + action.request.size
                    )
                    raise ValueError(
                        f"VisitAction '{action}' accesses memory region not fully loaded: {difference}"
                    )
                # Add to active visits
                active_visits.append(action)
            else:
                raise ValueError(f"Unknown action type: {type(action)}")

        # Concurrency check
        logger.info("Validating Actions - concurrency...")
        current_io_action = None
        final_io_action = None
        for action in self.actions[:-1]:
            if (
                current_io_action is not None
                and action.time >= current_io_action.time + current_io_action.cost()
            ):
                logger.debug(
                    f"--> io action '{current_io_action}' removed as completed"
                )
                current_io_action = None  # Reset on IO Action completion
            logger.debug(
                f"-> validating action: '{action}', time: {action.time}, cost: {action.cost()}"
            )
            if isinstance(action, (ReloadAction, OffloadAction)):
                if current_io_action is not None:
                    raise ValueError(
                        f"IO Actions '{current_io_action}' and '{action}' cannot be performed concurrently."
                    )
                current_io_action = action
                final_io_action = action
            elif isinstance(action, VisitAction):
                continue  # VisitAction can be concurrent with IO Actions
            else:
                raise ValueError(f"Unknown action type: {type(action)}")
        assert final_io_action is not None

        logger.debug("------------------------------")

        current_visit_action = None
        final_visit_action = None
        for start, reqs in self.requests_by_start.items():
            visit_actions_for_requests = [a for a in visit_actions if a.request in reqs]
            logger.debug(
                f"-> validating VisitActions for requests starting at {start}: {visit_actions_for_requests}"
            )
            assert len(visit_actions_for_requests) == len(reqs)
            longest_action = visit_actions_for_requests[0]
            for i in range(1, len(visit_actions_for_requests)):
                prev_action = visit_actions_for_requests[i - 1]
                curr_action = visit_actions_for_requests[i]
                if curr_action.time != prev_action.time:
                    raise ValueError(
                        f"VisitActions for requests ({reqs}) starting at {start} must be concurrent and start at the same time. \
'{curr_action}' and '{prev_action}' have different start times."
                    )
                if curr_action.cost() > longest_action.cost():
                    longest_action = curr_action
            if (
                current_visit_action is not None
                and longest_action.time
                < current_visit_action.time + current_visit_action.cost()
            ):
                raise ValueError(
                    f"VisitActions for requests ({reqs}) starting at {start} must be sequential. \
'{longest_action}' starts before previous VisitAction '{current_visit_action}' completes."
                )
            current_visit_action = longest_action
            final_visit_action = longest_action
        assert final_visit_action is not None

        # Report Correctness check
        logger.info("Validating Actions - report correctness...")
        last_action_end_time = max(
            final_io_action.time + final_io_action.cost(),
            final_visit_action.time + final_visit_action.cost(),
        )
        if self.actions[-1].time < last_action_end_time:
            raise ValueError(
                f"FinAction reports time {self.actions[-1].time}, but actual elapsed time is {last_action_end_time}"
            )
        elif self.actions[-1].time > last_action_end_time:
            logger.warning(
                f"FinAction reports time {self.actions[-1].time}, which is greater than actual elapsed time {last_action_end_time}. Consider reporting accurate time."
            )

        logger.info("Action Validated")
